# plot_value_vs_state.py

import json
import os
import pandas as pd
import plotly.express as px
import plotly.graph_objs as go


def load_trained_states(trained_states_file='trained_states.json'):
    """
    加载训练过程中记录的最优状态及其对应的值。
    """
    if os.path.exists(trained_states_file):
        with open(trained_states_file, 'r') as f:
            trained_states = json.load(f)
        print(f"Loaded {len(trained_states)} trained states from {trained_states_file}.")
        return trained_states
    else:
        print(f"No trained states found at {trained_states_file}.")
        return []


def validate_data(df):
    """
    验证DataFrame中的数据完整性和正确性。
    """
    required_columns = ['learning_rate', 'max_depth', 'n_estimators', 'value']
    if not set(required_columns).issubset(df.columns):
        print(f"DataFrame缺少必要的列。需要的列: {required_columns}")
        return False

    # 检查是否有缺失值
    if df[required_columns].isnull().any().any():
        print("DataFrame中存在缺失值。")
        return False

    print("Data validation passed.")
    return True


def create_scatter_plot(df, x_param, y_param, z_param=None, title="Value vs State Parameters"):
    """
    创建2D或3D散点图，展示value与超参数之间的关系。

    参数：
    - df (DataFrame): 包含超参数和value的数据框。
    - x_param (str): X轴对应的超参数。
    - y_param (str): Y轴对应的超参数。
    - z_param (str, optional): Z轴对应的超参数，用于3D散点图。
    - title (str): 图表标题。
    """
    if z_param:
        # 创建3D散点图
        fig = px.scatter_3d(
            df,
            x=x_param,
            y=y_param,
            z=z_param,
            color='value',
            size='value',
            hover_data=[x_param, y_param, z_param, 'value'],
            title=title,
            color_continuous_scale='Viridis'
        )
    else:
        # 创建2D散点图
        fig = px.scatter(
            df,
            x=x_param,
            y=y_param,
            color='value',
            size='value',
            hover_data=[x_param, y_param, 'value'],
            title=title,
            color_continuous_scale='Viridis'
        )

    fig.update_layout(
        coloraxis_colorbar=dict(
            title="Value"
        )
    )

    fig.show()

    # 保存为HTML文件（可选）
    # fig.write_html(f"{title.replace(' ', '_').lower()}.html")
    # print(f"Scatter plot saved as '{title.replace(' ', '_').lower()}.html'")


def create_bar_chart(df, x_param, title="Value of Best States"):
    """
    创建柱状图，展示不同超参数组合的Value。

    参数：
    - df (DataFrame): 包含超参数和value的数据框。
    - x_param (str): X轴对应的超参数列名，用于区分不同的条目。
    - title (str): 图表标题。
    """
    fig = px.bar(
        df,
        x=x_param,
        y='value',
        hover_data=[x_param, 'value'],
        title=title,
        labels={'value': 'Value', x_param: x_param.replace('_', ' ').title()}
    )

    fig.update_layout(
        yaxis=dict(range=[0, df['value'].max() * 1.1]),
        xaxis_tickangle=-45
    )

    fig.show()

    # 保存为HTML文件（可选）
    # fig.write_html(f"{title.replace(' ', '_').lower()}.html")
    # print(f"Bar chart saved as '{title.replace(' ', '_').lower()}.html'")


def create_line_plot(df, episode_col='Episode', value_col='value', title="Value Over Episodes"):
    """
    创建折线图，展示value随训练进程的变化。

    参数：
    - df (DataFrame): 包含episode和value的数据框。
    - episode_col (str): 代表训练进程的列名。
    - value_col (str): 代表value的列名。
    - title (str): 图表标题。
    """
    fig = px.line(
        df,
        x=episode_col,
        y=value_col,
        title=title,
        markers=True
    )

    fig.update_layout(
        yaxis=dict(title='Value'),
        xaxis=dict(title='Episode')
    )

    fig.show()

    # 保存为HTML文件（可选）
    # fig.write_html(f"{title.replace(' ', '_').lower()}.html")
    # print(f"Line plot saved as '{title.replace(' ', '_').lower()}.html'")


def main():
    # 加载最优训练状态
    trained_states = load_trained_states('trained_states_final.json')

    if not trained_states:
        print("No trained states to visualize. Exiting.")
        return

    # 转换为DataFrame
    df = pd.DataFrame(trained_states)

    # 数据验证
    if not validate_data(df):
        print("Data validation failed. Please check your data.")
        return

    # 创建2D散点图：Value vs Learning Rate
    create_scatter_plot(
        df,
        x_param='learning_rate',
        y_param='value',
        title="Value vs Learning Rate"
    )

    # 创建2D散点图：Value vs Max Depth
    create_scatter_plot(
        df,
        x_param='max_depth',
        y_param='value',
        title="Value vs Max Depth"
    )

    # 创建2D散点图：Value vs Number of Estimators
    create_scatter_plot(
        df,
        x_param='n_estimators',
        y_param='value',
        title="Value vs Number of Estimators"
    )

    # 创建3D散点图：Value vs Learning Rate, Max Depth, Number of Estimators
    create_scatter_plot(
        df,
        x_param='learning_rate',
        y_param='max_depth',
        z_param='n_estimators',
        title="3D Scatter Plot: Value vs Learning Rate, Max Depth, and Number of Estimators"
    )

    # 创建柱状图：不同 Learning Rate 的 Value
    create_bar_chart(
        df,
        x_param='learning_rate',
        title="Value of Best States by Learning Rate"
    )

    # 如果您的数据中包含Episode信息，并希望创建折线图，请确保DataFrame中有'Episode'列
    # 如果没有，可以添加一个基于DataFrame索引的Episode列
    if 'Episode' not in df.columns:
        df['Episode'] = range(1, len(df) + 1)

    # 创建折线图：Value随Episode变化
    create_line_plot(
        df,
        episode_col='Episode',
        value_col='value',
        title="Value Over Episodes"
    )

    # 您可以根据需要添加更多的可视化函数调用


if __name__ == '__main__':
    main()
